import json
import time
import subprocess

from absl import app, flags
from coin.tool.strat_monitor.util.info_util import _get_topics
from concurrent.futures import ProcessPoolExecutor, as_completed


FLAGS = flags.FLAGS


def _produce(topics):
  cmd = f"./pyrunner cc/experimental/wliang/ksqldb/ksqldb_producer.py  --kafka_topics=%s" % topics
  print(cmd)
  subprocess.run(cmd, shell=True)


def _consume(topics, ksqldb_server):
  cmd = f"./pyrunner cc/experimental/wliang/ksqldb/ksqldb_consumer.py  --kafka_topics=%s --ksqldb_server=%s" % (topics, ksqldb_server)
  print(cmd)
  subprocess.run(cmd, shell=True)


def _gen_topic_groups(topics, num_per_group):
  topic_groups = []
  topic_group = ""
  cnt = 0
  for topic in topics:
    topic_group = f"%s%s," % (topic_group, topic)
    cnt += 1
    if cnt == num_per_group:
      topic_group = topic_group if topic_group[-1] != ',' else topic_group[:-1]
      topic_groups.append(topic_group)
      topic_group = ""
      cnt = 0
  if len(topic_group) > 0:
    topic_group = topic_group if topic_group[-1] != ',' else topic_group[:-1]
    topic_groups.append(topic_group)
  return topic_groups

def main(_):
  topics = []
  topic_groups = [
      "STRAT_strategy-71.ap-northeast-1_model_smm_ftx_perpetual_0,STRAT_strategy-66.ap-northeast-1_basis_smm2_binance_perpetual_ftx2_perpetual,STRAT_strategy-155.ap-northeast-1_model_smm_ftx_perpetual_3,STRAT_strategy-41.ap-northeast-1_lm_volume_ftx_usd_swap_1,STRAT_strategy-15.ap-northeast-1_mm_portfolio_ftx_perpetual_v1"
  ]

  ksqldb_server_list = [f"http://localhost:%s" % x for x in FLAGS.ksqldb_server_ports.split(',')]
  ksqldb_server_count = len(ksqldb_server_list)

  #topics = _get_topics(None, '../coin_deploy/support_monitor/config/kafka_aws_config.json')
  #topic_groups = _gen_topic_groups(topics, 10)
  futures = []
  with ProcessPoolExecutor() as executor:
    idx = 1
    for topic_group in topic_groups:
      print("release group #%s %s" % (idx, topic_group))
      futures.append(executor.submit(_produce, topic_group))
      futures.append(executor.submit(_consume, topic_group, ksqldb_server_list[idx % ksqldb_server_count]))
      idx += 1
    print("release done. total of %s topics" % len(topics))
    for fut in as_completed(futures):
      _ = fut.result()


if __name__ == '__main__':
  flags.DEFINE_string('ksqldb_server_ports', "8088", "ksqldb server port list")
  app.run(main)
